import re
import uuid
import shlex
import shutil
from pathlib import Path
from typing import List, Dict, Tuple

import docker
from celery import Celery
from pymongo import MongoClient
from bson import ObjectId

# ───── Configuration ─────
SINGLE_CONTAINER = True           # False = 每用例新起容器
CPU_QUOTA        = 500_000_000    # 0.5 CPU
IMAGE_NAME       = "oj_judge:fixed"

# ───── Celery & Mongo ─────
celery_app = Celery(
    "judge",
    broker="redis://:hgzs.321@localhost:6379/0",
    backend="redis://:hgzs.321@localhost:6379/1"
)
db          = MongoClient("mongodb://localhost:27017").oj
subs, probs = db.submissions, db.problems
docker_cli  = docker.from_env()

# ───── Helpers ─────
strip_ws  = lambda s: re.sub(r"\s+", "", s)
TIME_RE   = re.compile(r"^\s*([\d.]+)\s+(\d+)\s*$", re.M)
severity  = {
    "Accepted":0, "Presentation Error":1, "Wrong Answer":2,
    "Output Limit Exceeded":3, "Runtime Error":4,
    "Memory Limit Exceeded":5, "Time Limit Exceeded":6
}

def mem2bytes(s: str) -> int:
    s = s.lower()
    if s.endswith("g"): return int(float(s[:-1]) * 1024**3)
    if s.endswith("m"): return int(float(s[:-1]) * 1024**2)
    if s.endswith("k"): return int(float(s[:-1]) * 1024)
    return int(float(s))

def first_diff(a: str, b: str) -> int:
    for i, (x, y) in enumerate(zip(a, b)):
        if x != y:
            return i
    return min(len(a), len(b))

def safe_peak_kb(mem_stats: dict) -> int:
    mem_info = mem_stats.get("memory_stats", {})
    return (
        mem_info.get("max_usage")
        or mem_info.get("stats", {}).get("max_usage")
        or mem_info.get("usage", 0)
    ) // 1024

# ───── Container helpers ─────
def start_base_container(work: Path, mem_k: int):
    return docker_cli.containers.run(
        IMAGE_NAME, "sleep infinity",
        volumes={str(work): {'bind': '/home/runner/sandbox', 'mode': 'rw'}},
        working_dir="/home/runner/sandbox",
        network_disabled=True,
        mem_limit=f"{mem_k}k",
        nano_cpus=CPU_QUOTA,
        detach=True, tty=False, stdin_open=False
    )

def exec_cmd(ctn, cmd: str, tl: int) -> Tuple[int, str, str, int, int]:
    wrapped = f"timeout {tl}s /usr/bin/time -f '%e %M' bash -lc {shlex.quote(cmd)}"
    eid = docker_cli.api.exec_create(ctn.id, wrapped)["Id"]
    out = docker_cli.api.exec_start(eid, demux=True)
    info = docker_cli.api.exec_inspect(eid)
    stdout = (out[0] or b"").decode("utf-8", "ignore")
    stderr = (out[1] or b"").decode("utf-8", "ignore")
    m = TIME_RE.search(stderr)
    run_ms  = int(float(m.group(1)) * 1000) if m else 0
    peak_kb = int(m.group(2))             if m else 0
    return info["ExitCode"], stdout.strip(), stderr.strip(), run_ms, peak_kb

# ───── Celery task ─────
@celery_app.task(name="judge_code")
def judge_code(sub_id: str):
    # 标记为编译中
    subs.update_one({"_id": ObjectId(sub_id)}, {"$set": {"status": "Compiling"}})
    sub  = subs.find_one({"_id": ObjectId(sub_id)})
    prob = probs.find_one({"_id": sub["problem_id"]})

    # 创建临时工作目录并写入代码
    work = Path(f"/tmp/{uuid.uuid4().hex}")
    work.mkdir(exist_ok=True)
    src = work / ("main.py" if sub["lang"] == "py" else "main.cpp")
    src.write_text(sub["code"], encoding="utf-8")

    # --- 编译 C++ ---
    if sub["lang"] == "cpp":
        try:
            docker_cli.containers.run(
                IMAGE_NAME, "g++ main.cpp -O2 -std=c++17 -o a.out",
                volumes={str(work): {'bind': '/home/runner/sandbox', 'mode': 'rw'}},
                working_dir="/home/runner/sandbox",
                network_disabled=True,
                mem_limit="512m", nano_cpus=CPU_QUOTA,
                remove=True, stdout=True, stderr=True
            )
        except docker.errors.ContainerError as e:
            subs.update_one(
                {"_id": sub["_id"]},
                {"$set": {
                    "status": "Compile Error",
                    "compile_error": e.stderr.decode('utf-8', 'ignore')
                }}
            )
            shutil.rmtree(work, ignore_errors=True)
            return

    # 计算资源限制
    mem_k = mem2bytes(f"{prob['memory_limit']}m") // 1024
    tl_s  = max(1, prob["time_limit"] // 1000)
    base  = start_base_container(work, mem_k) if SINGLE_CONTAINER else None

    # 如果复用容器，预先测一次空命令的开销
    overhead_ms = 0
    if SINGLE_CONTAINER:
        _, _, _, overhead_ms, _ = exec_cmd(base, "true", tl_s)

    subs.update_one({"_id": sub["_id"]}, {"$set": {"status": "Running & Judging"}})
    cases = []
    final = "Accepted"

    for idx, tc in enumerate(prob["tests"], start=1):
        runner = "python3 -u main.py" if sub["lang"] == "py" else "./a.out"
        cmd    = f"printf '%s\\n' {shlex.quote(tc['in'].rstrip())} | {runner}"

        if SINGLE_CONTAINER:
            code, out, err, measured_ms, _ = exec_cmd(base, cmd, tl_s)
            run_ms = max(0, measured_ms - overhead_ms)
            base.reload()
            oom = base.attrs["State"]["OOMKilled"]
            mem_stats = docker_cli.api.stats(base.id, stream=False)
            real_peak = safe_peak_kb(mem_stats)
        else:
            full = f"timeout {tl_s}s /usr/bin/time -f '%e %M' bash -lc {shlex.quote(cmd)}"
            ctn  = docker_cli.containers.run(
                IMAGE_NAME, full,
                volumes={str(work): {'bind': '/home/runner/sandbox', 'mode': 'rw'}},
                working_dir="/home/runner/sandbox",
                network_disabled=True,
                mem_limit=f"{mem_k}k",
                nano_cpus=CPU_QUOTA,
                detach=True, tty=False, stdin_open=False
            )
            info = ctn.wait()
            code = info["StatusCode"]
            out  = ctn.logs(stdout=True, stderr=False).decode("utf-8", "ignore").strip()
            err  = ctn.logs(stdout=False, stderr=True).decode("utf-8", "ignore")
            ctn.reload()
            oom = ctn.attrs["State"]["OOMKilled"]
            mem_stats = ctn.stats(stream=False)
            real_peak = safe_peak_kb(mem_stats)
            m = TIME_RE.search(err)
            run_ms     = int(float(m.group(1)) * 1000) if m else 0
            peak_kb    = int(m.group(2))             if m else 0
            ctn.remove(force=True)

        # --- 判题逻辑 ---
        if oom or real_peak > mem_k:
            verdict = "Memory Limit Exceeded"
        elif code == 124:
            verdict = "Time Limit Exceeded"
        elif code != 0:
            verdict = "Runtime Error"
        else:
            expect, output = tc["out"].strip(), out
            if len(output) > 2 * len(expect):
                verdict = "Output Limit Exceeded"
            elif output != expect:
                verdict = ("Presentation Error" if strip_ws(output) == strip_ws(expect)
                           else "Wrong Answer")
            else:
                verdict = "Accepted"

        rec = {"case": idx, "verdict": verdict, "time": run_ms, "mem": real_peak}
        if verdict in ("Wrong Answer", "Presentation Error"):
            rec.update({
                "output": output,
                "expect": expect,
                "first_mismatch_pos": first_diff(output, expect)
            })
        if verdict.startswith("Runtime"):
            rec["error"] = err

        cases.append(rec)
        if severity[verdict] > severity[final]:
            final = verdict

    if base:
        base.remove(force=True)

    subs.update_one(
        {"_id": sub["_id"]},
        {"$set": {"status": final, "cases": cases}}
    )
    shutil.rmtree(work, ignore_errors=True)
